import uuid
from pathlib import Path

import ffmpeg
from pydantic import BaseModel
from loguru import logger

def query_video_info(file_path: Path):
    probe = ffmpeg.probe(file_path)

    ret = MediaInfo()

    for stream in probe['streams']:
        if stream['codec_type'] == 'video':
            ret.video_streams.append(VideoStreamInfo(**stream))
        elif stream['codec_type'] == 'audio':
            ret.audio_streams.append(AudioStreamInfo(**stream))

    return ret


def separate_stream(file_path: Path, audio_params: dict = None , video_params: dict = None) -> dict[str, Path]:
    """
    分离视频文件的视频流与音频流
    :param file_path:
    :param audio_params: 视频编码参数
    :param video_params: 视频编码参数
    :return: 以video_stream和audio_stream为key的dict，value为对应的文件路径
    """
    probe = ffmpeg.probe(file_path)
    if len(probe['streams']) == 0:
        raise ValueError("No streams found in the video file")

    has_handle_video_stream, has_handle_audio_stream = False, False
    video_stream, audio_stream = None, None

    for stream in probe['streams']:
        if stream['codec_type'] == 'video' and not has_handle_video_stream:
            video_stream = _separate_video_stream(file_path, video_params)
        elif stream['codec_type'] == 'audio' and not has_handle_audio_stream:
            audio_stream = _separate_audio_stream(file_path, audio_params)

    return {
        "video_stream": video_stream,
        "audio_stream": audio_stream
    }


def extract_frame_by_fps(video_file_path: Path, target_dir: Path, fps: float = 0.5):
    """
    按照指定的fps对视频文件生成截图
    :param video_file_path: 源文件路径
    :param target_dir: 生成的截图文件存放的文件夹的路径
    :param fps: 每秒帧数
    :return: 按照时间顺序产生的截图文件的路径
    """

    if not video_file_path.is_file():
        raise ValueError("video_file_path is not a file")

    if not target_dir.is_dir():
        raise ValueError("target_dir is not a dir")

    input = ffmpeg.input(str(video_file_path))
    uuid_str = str(uuid.uuid4())
    ffmpeg.output(input, str(target_dir) + f"/{uuid_str}_%d.jpg", vf=f"fps={fps}").run()

    frame_img_list = list(target_dir.glob(f"{uuid_str}_*.jpg"))
    frame_img_list.sort(key = lambda x: int(x.name.split(".")[0].split("_")[1]))

    return frame_img_list


def audio_transcode(source: Path, target: Path, **kwargs):
    """
    对音频文件进行转码
    :param source: 源文件路径
    :param target: 目标文件路径
    :param kwargs: 转码参数
        * sample_rate: 采样率
        * channels: 声道数
        * sample_format: 采样格式

    :return:
    """
    if len(kwargs) == 0:
        raise ValueError("transcode params required")

    if target is None:
        target = source.parent / (str(source.name) + "_transcode")

    input = ffmpeg.input(str(source))

    params = {}
    for key, value in kwargs.items():
        match key:
            case 'sample_rate':
                params["ar"] = value
            case 'channels':
                params["ac"] = value
            case 'sample_format':
                params["sample_fmt"] = value
            case _:
                logger.warning(f"received unknow args <{key}>")

    if len(params) == 0:
        raise ValueError("no valid transcode params")

    ffmpeg.output(input, str(target), **params).run()
    return target


def _separate_video_stream(file_path: Path, params: dict):
    """
    分离视频文件的视频流
    :param file_path:
    :param params: 视频编码参数
    :return:
    """
    if params is None:
        params = {'map': '0:v', 'c': 'copy', 'f': 'mp4'}

    video_stream_separate_result = file_path.parent / (str(file_path.name) + "_video")

    input = ffmpeg.input(str(file_path))
    ffmpeg.output(input, str(video_stream_separate_result), **params).run()

    return video_stream_separate_result


def _separate_audio_stream(file_path: Path, params: dict):
    """
    分离视频文件的音频流
    :param file_path:
    :param params: 音频编码参数
    :return:
    """

    audio_stream_separate_result = file_path.parent / (str(file_path.name) + "_audio")

    if params is None:
        params = {'c:a': 'libmp3lame', 'f': 'mp3'}

    input = ffmpeg.input(str(file_path))

    ffmpeg.output(input, str(audio_stream_separate_result), **params).run()

    return audio_stream_separate_result



class VideoStreamInfo(BaseModel):
    """
    视频流信息
    """
    """
    流编号
    """
    index:  int
    """
    编码名称
    """
    codec_name:  str
    """
    编码名称完整版
    """
    codec_long_name:  str
    """
    编码tag(应该就是封装容器类型名)
    """
    codec_tag_string:  str
    """
    水平分辨率
    """
    width:  int
    """
    垂直分辨率
    """
    height:  int
    """
    编码水平分辨率
    """
    coded_width:  int
    """
    编码垂直分辨率
    """
    coded_height:  int
    """
    像素格式（色彩采样格式）
    """
    pix_fmt:  str
    """
    时间基
    """
    time_base:  str
    """
    开始时间戳
    """
    start_pts: int
    """
    开始时间
    """
    start_time:  float
    """
    时长时间戳
    """
    duration_ts: int
    """
    时长
    """
    duration:  float
    """
    码率(bps)
    """
    bit_rate: int

class AudioStreamInfo(BaseModel):
    """
    音频流信息
    """
    """
    流编号
    """
    index:  int
    """
    编码名称
    """
    codec_name: str
    """
    编码名称完整版
    """
    codec_long_name: str
    """
    编码tag(应该就是封装容器类型名)
    """
    codec_tag_string: str
    """
    采样格式
    """
    sample_fmt: str
    """
    采样率
    """
    sample_rate: int
    """
    声道数
    """
    channels: int
    """
    声道布局
    """
    channel_layout: str
    """
    时间基
    """
    time_base:  str
    """
    开始时间戳
    """
    start_pts: int
    """
    开始时间
    """
    start_time:  float
    """
    时长时间戳
    """
    duration_ts: int
    """
    时长
    """
    duration:  float
    """
    码率(bps)
    """
    bit_rate: int




class MediaInfo(BaseModel):
    """
    媒体文件信息
    """

    video_streams: list[VideoStreamInfo] = []
    audio_streams: list[AudioStreamInfo] = []

